"""
model name : 深度学习
file       : oloss.py
information:
    author : OuYang
    time   : 2025/1/22
"""

import torch
from torch import nn


class YOLOLoss(nn.Module):
    def __init__(self, coord=5, noobj=.5, s=7, b=2):
        super(YOLOLoss, self).__init__()
        self.coord = coord
        self.noobj = noobj
        self.s = s
        self.b = b
        self.cell_size = 1.0 / self.s

    def forward(self, pred, target):

        cls_loss = torch.tensor(0, dtype=torch.float32, device=target.device)
        conf_loss = torch.tensor(0, dtype=torch.float32, device=target.device)
        location_loss = torch.tensor(0, dtype=torch.float32, device=target.device)
        for batch_idx in range(pred.size(0)):
            for i in range(pred.size(1)):
                for j in range(pred.size(2)):
                    ious = []
                    max_iou = -1
                    max_bbox_idx = 0
                    for bbox_idx in range(self.b):
                        pred_x = pred[batch_idx, i, j, bbox_idx * 5]
                        pred_y = pred[batch_idx, i, j, bbox_idx * 5 + 1]
                        pred_w = pred[batch_idx, i, j, bbox_idx * 5 + 2]
                        pred_h = pred[batch_idx, i, j, bbox_idx * 5 + 3]

                        target_x = target[batch_idx, i, j, bbox_idx * 5]
                        target_y = target[batch_idx, i, j, bbox_idx * 5]
                        target_w = target[batch_idx, i, j, bbox_idx * 5]
                        target_h = target[batch_idx, i, j, bbox_idx * 5]

                        iou = self._iou(
                            box1=(pred_x, pred_y, pred_w, pred_h),
                            box2=(target_x, target_y, target_w, target_h),
                            i=i, j=j
                        )
                        ious.append(iou)

                        if iou > max_iou:
                            max_iou = iou
                            max_bbox_idx = bbox_idx

                    if target[batch_idx, i, j, 4] == 1:
                        # cls loss
                        cls_loss += torch.sum(
                            ((pred[batch_idx, i, j, self.b * 5:] - target[batch_idx, i, j, self.b * 5:]) ** 2)
                        )

                        # location loss
                        location_loss += self.coord * torch.sum(
                            (pred[batch_idx, i, j, max_bbox_idx * 5] - target[batch_idx, i, j, max_bbox_idx]) ** 2 +
                            (pred[batch_idx, i, j, max_bbox_idx * 5 + 1] - target[
                                batch_idx, i, j, max_bbox_idx + 1]) ** 2
                        )
                        location_loss += self.coord * torch.sum(
                            (pred[batch_idx, i, j, max_bbox_idx * 5 + 2].sqrt() - target[
                                batch_idx, i, j, max_bbox_idx * 5 + 2].sqrt()) ** 2 +
                            (pred[batch_idx, i, j, max_bbox_idx * 5 + 3].sqrt() - target[
                                batch_idx, i, j, max_bbox_idx * 5 + 3].sqrt()) ** 2
                        )

                        # conf_loss
                        for iou_idx, iou in enumerate(ious):
                            if iou_idx == max_bbox_idx:
                                conf_loss += torch.sum(
                                    (iou - pred[batch_idx, i, j, max_bbox_idx]) ** 2
                                )
                            else:
                                conf_loss += self.noobj * torch.sum(
                                    (iou - pred[batch_idx, i, j, max_bbox_idx]) ** 2
                                )

        total_loss = conf_loss + cls_loss + location_loss
        return total_loss, cls_loss, conf_loss, location_loss

    def _iou(self, box1, box2, i, j):
        b1x, b1y, b1w, b1h = box1
        b2x, b2y, b2w, b2h = box2

        b1x1 = b1x * self.cell_size + i * self.cell_size - b1w / 2
        b1y1 = b1y * self.cell_size + j * self.cell_size - b1h / 2
        b1x2 = b1x * self.cell_size + i * self.cell_size + b1w / 2
        b1y2 = b1y * self.cell_size + j * self.cell_size + b1h / 2

        b2x1 = b2x * self.cell_size + i * self.cell_size - b2w / 2
        b2y1 = b2y * self.cell_size + j * self.cell_size - b2h / 2
        b2x2 = b2x * self.cell_size + i * self.cell_size + b2w / 2
        b2y2 = b2y * self.cell_size + j * self.cell_size + b2h / 2

        if b1x2 < b2x1 or b1x1 > b2x2 or b1y2 < b2y1 or b1y1 > b2y2:
            return 0
        inter_x1 = torch.max(b1x1, b2x1)
        inter_y1 = torch.max(b1y1, b2y1)
        inter_x2 = torch.min(b1x2, b2x2)
        inter_y2 = torch.min(b1y2, b2y2)

        b1_area = b1w * b1h
        b2_area = b2w * b2h
        inter_area = torch.abs(inter_x2 - inter_x1) * torch.abs(inter_y2 - inter_y1)

        union_area = b1_area + b2_area - inter_area
        iou = inter_area / torch.clamp(union_area, 1e-6)

        return iou if inter_area > 0 else 0


# if __name__ == '__main__':
#     output = torch.tensor([
#         [[0.3, 0.4, 0.5, 0.6, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.2, 0.1, 0.5],
#          [0.1, 0.3, 0.4, 0.2, 0.2, 0.4, 0.2, 0.2, 0.4, 0.3, 0.3, 0.4, 0.2]],
#         [[0.1, 0.2, 0.2, 0.2, 0.3, 0.4, 0.3, 0.3, 0.3, 0.2, 0.4, 0.3, 0.2],
#          [0.1, 0.1, 0.1, 0.1, 0.8, 0.5, 0.4, 0.3, 0.2, 0.1, 0.5, 0.3, 0.2]]
#     ], dtype=torch.float32)
#     output = torch.reshape(output, (-1, 2, 2, 13))
#
#     target = torch.tensor([
#         [[0.2, 0.3, 0.4, 0.5, 1.0, 0.2, 0.3, 0.4, 0.5, 1, 0., 0., 1],
#          [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]],
#         [[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
#          [0.2, 0.1, 0.3, 0.2, 1.0, 0.2, 0.1, 0.3, 0.2, 1, 1.0, 0., 0.]]
#     ], dtype=torch.float32)
#     target = torch.reshape(target, (-1, 2, 2, 13))
#
#     print(output.shape)
#     print(target.shape)
#
#     loss_function = YOLOLoss(b=2, s=2)
#     loss = loss_function(output, target)
#     print(loss)
if __name__ == '__main__':
    output = torch.tensor([
        [[[0.4418, 0.7047, 0.3439, 0.6775, 0.7019, 0.5227, 0.6190, 0.3534,
           0.2684, 0.8168, 0.6451],
          [0.3130, 0.3774, 0.2232, 0.2649, 0.6049, 0.4615, 0.7493, 0.2419,
           0.5396, 0.3336, 0.4666],
          [0.6307, 0.6902, 0.4919, 0.3114, 0.8004, 0.5673, 0.2284, 0.6419,
           0.1597, 0.6273, 0.4472],
          [0.5652, 0.1981, 0.4890, 0.6149, 0.4008, 0.3854, 0.7306, 0.3327,
           0.4405, 0.3843, 0.6046],
          [0.5760, 0.6845, 0.5472, 0.5386, 0.7579, 0.4034, 0.2187, 0.6795,
           0.7319, 0.6481, 0.4207],
          [0.7104, 0.1309, 0.7795, 0.3454, 0.7786, 0.4297, 0.2523, 0.7487,
           0.8688, 0.4029, 0.4904],
          [0.4463, 0.7735, 0.5358, 0.2501, 0.2933, 0.6874, 0.5515, 0.5933,
           0.1305, 0.5640, 0.2501]],

         [[0.6008, 0.5906, 0.5964, 0.7702, 0.7443, 0.5181, 0.1830, 0.6916,
           0.7833, 0.4761, 0.1403],
          [0.0918, 0.1100, 0.2696, 0.4586, 0.0958, 0.4719, 0.5124, 0.2499,
           0.4554, 0.8776, 0.7245],
          [0.5283, 0.6462, 0.1235, 0.6089, 0.5752, 0.8205, 0.7489, 0.0424,
           0.6722, 0.6518, 0.6941],
          [0.3561, 0.6586, 0.7465, 0.6396, 0.6116, 0.5713, 0.6284, 0.1908,
           0.7726, 0.3797, 0.3772],
          [0.5256, 0.1152, 0.2378, 0.3420, 0.8400, 0.4943, 0.6770, 0.2809,
           0.1306, 0.4810, 0.2678],
          [0.7993, 0.8461, 0.3608, 0.4667, 0.3248, 0.8768, 0.3222, 0.2610,
           0.4860, 0.3922, 0.3275],
          [0.6045, 0.8918, 0.8122, 0.3322, 0.1821, 0.5455, 0.5385, 0.5806,
           0.6700, 0.3083, 0.1561]],

         [[0.7209, 0.6274, 0.2959, 0.7353, 0.3458, 0.7711, 0.1770, 0.4425,
           0.1910, 0.7297, 0.5000],
          [0.5285, 0.3965, 0.2690, 0.5949, 0.2757, 0.6371, 0.5112, 0.4660,
           0.5060, 0.2532, 0.2156],
          [0.4776, 0.5530, 0.6740, 0.2474, 0.3049, 0.3356, 0.6500, 0.1623,
           0.2121, 0.6396, 0.4114],
          [0.6590, 0.6353, 0.5328, 0.5326, 0.5396, 0.4671, 0.1567, 0.5079,
           0.4372, 0.4873, 0.7350],
          [0.8625, 0.8086, 0.2258, 0.4895, 0.4460, 0.8230, 0.5941, 0.7246,
           0.4083, 0.5240, 0.6491],
          [0.4620, 0.4378, 0.2287, 0.5379, 0.8319, 0.3350, 0.8280, 0.7114,
           0.5393, 0.5117, 0.6558],
          [0.6841, 0.8469, 0.1391, 0.6996, 0.4249, 0.1108, 0.4733, 0.3351,
           0.1077, 0.5843, 0.2145]],

         [[0.5697, 0.1620, 0.2940, 0.2100, 0.7556, 0.3123, 0.8239, 0.5134,
           0.1866, 0.8338, 0.8512],
          [0.4366, 0.2782, 0.5830, 0.4168, 0.1671, 0.6618, 0.6229, 0.5080,
           0.7862, 0.3052, 0.5349],
          [0.2686, 0.5325, 0.6727, 0.3665, 0.5521, 0.7650, 0.2797, 0.3869,
           0.6601, 0.6303, 0.7181],
          [0.6815, 0.5689, 0.9352, 0.2175, 0.9398, 0.6083, 0.6644, 0.6994,
           0.5610, 0.7085, 0.6309],
          [0.7244, 0.3706, 0.8627, 0.5563, 0.4444, 0.4487, 0.8211, 0.6925,
           0.5472, 0.2073, 0.4491],
          [0.7213, 0.5865, 0.2936, 0.2519, 0.5225, 0.6264, 0.3243, 0.3843,
           0.8439, 0.8522, 0.8142],
          [0.7587, 0.5791, 0.5316, 0.2631, 0.2782, 0.5173, 0.3375, 0.4251,
           0.6108, 0.5247, 0.3470]],

         [[0.5593, 0.4925, 0.8682, 0.5752, 0.5148, 0.2929, 0.6005, 0.2957,
           0.6661, 0.3208, 0.2088],
          [0.2091, 0.5091, 0.7278, 0.6247, 0.1752, 0.3890, 0.2855, 0.7533,
           0.8032, 0.3039, 0.3799],
          [0.4229, 0.3375, 0.2178, 0.0757, 0.4826, 0.0674, 0.1207, 0.5152,
           0.2519, 0.6862, 0.4886],
          [0.3257, 0.5053, 0.8154, 0.4205, 0.4098, 0.4067, 0.4000, 0.8341,
           0.6047, 0.4316, 0.2672],
          [0.5575, 0.6941, 0.5996, 0.7116, 0.2102, 0.7792, 0.5458, 0.7273,
           0.3767, 0.2876, 0.6472],
          [0.2044, 0.7437, 0.5215, 0.4009, 0.9483, 0.4849, 0.3984, 0.4049,
           0.5900, 0.7372, 0.8579],
          [0.6325, 0.3327, 0.5054, 0.5586, 0.6356, 0.4075, 0.5652, 0.2124,
           0.3480, 0.6055, 0.5149]],

         [[0.3572, 0.3065, 0.3576, 0.2716, 0.8078, 0.2872, 0.7204, 0.2267,
           0.2840, 0.1842, 0.6847],
          [0.7023, 0.5524, 0.6493, 0.3383, 0.3768, 0.4475, 0.5983, 0.2499,
           0.7079, 0.5822, 0.8458],
          [0.7038, 0.5061, 0.3639, 0.3474, 0.1599, 0.3628, 0.3696, 0.8733,
           0.7290, 0.3470, 0.5815],
          [0.4647, 0.2058, 0.4215, 0.2925, 0.7337, 0.7470, 0.2515, 0.4313,
           0.2600, 0.3362, 0.2324],
          [0.3875, 0.2521, 0.3422, 0.6496, 0.3927, 0.7192, 0.5011, 0.4479,
           0.8063, 0.4671, 0.5043],
          [0.4121, 0.4507, 0.4426, 0.4422, 0.3348, 0.4272, 0.8004, 0.6412,
           0.5106, 0.4172, 0.5113],
          [0.2189, 0.7003, 0.6216, 0.1922, 0.3731, 0.4973, 0.5769, 0.2652,
           0.0983, 0.2784, 0.1779]],

         [[0.5239, 0.3609, 0.3708, 0.2389, 0.2443, 0.7240, 0.2897, 0.2315,
           0.5909, 0.7604, 0.1223],
          [0.7298, 0.3849, 0.6499, 0.4313, 0.6393, 0.6097, 0.2284, 0.3771,
           0.7688, 0.5399, 0.6899],
          [0.5238, 0.4720, 0.5599, 0.2417, 0.8726, 0.6068, 0.1922, 0.5322,
           0.4356, 0.4069, 0.4213],
          [0.2991, 0.2504, 0.6652, 0.8927, 0.2486, 0.5034, 0.4908, 0.6466,
           0.8062, 0.3496, 0.5773],
          [0.2402, 0.4491, 0.7078, 0.3222, 0.6686, 0.4601, 0.3871, 0.4238,
           0.6649, 0.6160, 0.7407],
          [0.4464, 0.5330, 0.2347, 0.7336, 0.7189, 0.8564, 0.2438, 0.8681,
           0.3638, 0.2079, 0.2233],
          [0.7770, 0.5294, 0.7180, 0.8786, 0.8545, 0.5212, 0.5873, 0.4132,
           0.5661, 0.6178, 0.6553]]]
    ], dtype=torch.float32)
    target = torch.tensor([
        [[[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000]],

         [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000]],

         [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000]],

         [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.6203, 0.3688, 0.4500, 0.3219, 1.0000, 0.6203, 0.3688, 0.4500,
           0.3219, 1.0000, 1.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000]],

         [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000]],

         [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000]],

         [[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000],
          [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
           0.0000, 0.0000, 0.0000]]]
    ], dtype=torch.float32)
    loss_fn = YOLOLoss(b=2)
    loss = loss_fn(output, target)

    print(loss)
